import java.text.DecimalFormat;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.text.*;
import java.util.*;
import java.io.*;

public class HMM extends MarkovModel {

  protected static int maxIterationsBaumWelsh = 250;
  protected static boolean debugBaumWelsh = true;

  protected int nstate;
  protected Object[] state;
  protected double[][] loga;

  protected double[][] loge;
    
  public HMM(Object[] state, double[][] amat, Object[] esym, double[][] emat ) {
    if (state.length != amat.length) throw new IllegalArgumentException("HMM: state and amat disagree");
    if (amat.length != emat.length) throw new IllegalArgumentException("HMM: amat and emat disagree");
    for (int i=0; i<amat.length; i++) {
      if (state.length != amat[i].length) throw new IllegalArgumentException("HMM: amat non-square");
      if (esym.length != emat[i].length) throw new IllegalArgumentException("HMM: esym and emat disagree");
    }
    nstate = state.length + 1;
    this.state = new Object[nstate];
    loga = new double[nstate][nstate];
    this.state[0] = (Object)(new String("---INITIAL-STATE---"));
    loga[0][0] = Double.NEGATIVE_INFINITY;
    double fromstart = Math.log(1.0/state.length);
    for (int j=1; j<nstate; j++) loga[0][j] = fromstart;
    for (int i=1; i<nstate; i++) {
      this.state[i] = state[i-1];
      loga[i][0] = Double.NEGATIVE_INFINITY;
      for (int j=1; j<nstate; j++) loga[i][j] = Math.log(amat[i-1][j-1]);
    }
    this.esym = esym;
    nesym = esym.length;
    loge = new double[emat.length + 1][esym.length];
    for (int b=0; b < nesym; b++) {
      loge[0][b] = Double.NEGATIVE_INFINITY;
      for (int k=0; k<emat.length; k++) loge[k+1][b] = Math.log(emat[k][b]);
    }
  }

  public void print(PrintStream out) { printa(out); printe(out); }

  public void printa(PrintStream out) {
    out.println("Transition probabilities:");
    for (int i=1; i<nstate; i++) {
      for (int j=1; j<nstate; j++) out.print(fmtlog(loga[i][j]));
      out.println();
    }
  }

  public void printe(PrintStream out) {
    out.println("Emission probabilities:");
    for (int b=0; b<nesym; b++) out.print(esym[b] + hdrpad);
    out.println();
    for (int i=1; i<loge.length; i++) {
      for (int b=0; b<nesym; b++) out.print(fmtlog(loge[i][b]));
      out.println();
    }
  }

  // Input.
  // xs    is the set of training sequences
  // state is the set of HMM state names
  // esym  is the set of emissible symbols
  public static HMM baumwelch(Object[][] xs, Object[] state, Object[] esym, final double threshold ) {
    int nstate = state.length;
    int nseqs  = xs.length;
    int nesym  = esym.length;
    Forward[] fwds = new Forward[nseqs];
    Backward[] bwds = new Backward[nseqs];
    double[] logP = new double[nseqs];
    double[][] amat = new double[nstate][];
    double[][] emat = new double[nstate][];
    int[] esyminv = new int[esym.length];
    for (int i=0; i<esyminv.length; i++) esyminv[i] = -1;
    for (int b=0; b<nesym; b++) esyminv[b] = b;
    for (int k=0; k<nstate; k++) {
      amat[k] = randomdiscrete(nstate);
      emat[k] = randomdiscrete(nesym);
    }
    HMM hmm = new HMM(state, amat, esym, emat);  
    double oldloglikelihood;
    double loglikelihood = fwdbwd(hmm, xs, fwds, bwds, logP);
    int numIterations = 0;
    do {
      if ( maxIterationsBaumWelsh > -1 && ( numIterations++ > maxIterationsBaumWelsh ) ) break;
      oldloglikelihood = loglikelihood;
      double[][] A = new double[nstate][nstate];
      double[][] E = new double[nstate][nesym];
      for (int s=0; s < nseqs; s++) {
        Object[] x = xs[s];
        Forward fwd  = fwds[s];
        Backward bwd = bwds[s];
        int L = x.length;
        double P = logP[s];
        for (int i=0; i < L; i++) {
          for (int k=0; k<nstate; k++) 
            E[k][esyminv[hmm.getSym(x[i])]] += exp(fwd.f[i+1][k+1] + bwd.b[i+1][k+1] - P);
        }
        for (int i=0; i<L-1; i++) 
          for (int k=0; k<nstate; k++) 
            for (int ell=0; ell<nstate; ell++) 
              A[k][ell] += exp(fwd.f[i+1][k+1] 
                               + hmm.loga[k+1][ell+1] 
                               + hmm.loge[ell+1][hmm.getSym(x[i+1])] 
                               + bwd.b[i+2][ell+1] 
                               - P);
      }
      for (int k=0; k<nstate; k++) {
        double Aksum = 0;
        for (int ell=0; ell<nstate; ell++) Aksum += A[k][ell];
        for (int ell=0; ell<nstate; ell++) amat[k][ell] = ( A[k][ell] + 1.0 ) / ( Aksum + 0.0 + nstate );
        double Eksum = 0;
        for (int b=0; b<nesym; b++) Eksum += E[k][b];
        for (int b=0; b<nesym; b++) emat[k][b] = ( E[k][b] + 1.0 ) / ( Eksum + 0.0 + nesym);
      }
      hmm = new HMM(state, amat, esym, emat);
      loglikelihood = fwdbwd(hmm, xs, fwds, bwds, logP);
      if ( debugBaumWelsh ) System.err.println("Baum-Welch Training - log likelihood = " + loglikelihood + " ( diff from previous is " + Math.abs(oldloglikelihood - loglikelihood) + ")");
    } while (Math.abs(oldloglikelihood - loglikelihood) > threshold);
   
    return hmm;
  }

  protected static double fwdbwd(HMM hmm, Object[][] xs, Forward[] fwds, Backward[] bwds, double[] logP) {
    double loglikelihood = 0;
    for (int s=0; s<xs.length; s++) {
      fwds[s] = new Forward(hmm, xs[s]);
      bwds[s] = new Backward(hmm, xs[s]);
      logP[s] = fwds[s].logprob();
      loglikelihood += logP[s];
    }
    return loglikelihood;
  }

  public static double exp(double x) {
    if (x == Double.NEGATIVE_INFINITY) return 0;
    else return Math.exp(x);
  }

  protected static double[] uniformdiscrete(int n) {
    double[] ps = new double[n];
    for (int i=0; i<n; i++) ps[i] = 1.0/n;
    return ps;
  }    

  protected static double[] randomdiscrete(int n) {
    double[] ps = new double[n];
    double sum = 0;
    for (int i=0; i<n; i++) {
      ps[i] = rnd.nextDouble();
      sum += ps[i];
    }
    for (int i=0; i<n; i++) ps[i] /= sum;
    return ps;
  }

}
